import cv2
import tkinter as tk
from tkinter import filedialog
from ultralytics import YOLO
import time
import numpy as np
from datetime import datetime




class FaceDetectionApp:
    def __init__(self, root):
        self.root = root
        self.root.title("YOLO Face Detection")

        # YOLO模型选择
        self.model = None
        self.model_label = tk.Label(self.root, text="选择YOLO模型:")
        self.model_label.pack(pady=5)

        self.model_path_button = tk.Button(self.root, text="选择YOLO模型路径", command=self.choose_model_path)
        self.model_path_button.pack(pady=5)
        self.model_path = None

        # 路径选择
        self.path_button = tk.Button(self.root, text="选择保存路径", command=self.choose_save_path)
        self.path_button.pack(pady=5)
        self.save_path = None

        # 开始检测按钮
        self.start_button = tk.Button(self.root, text="开始检测", command=self.start_detection)
        self.start_button.pack(pady=20)

        # 初始化摄像头和其他变量
        self.prev_frame = None
        self.face_detected = False
        self.frame_rate = 5  # 帧率，每5秒检测一次变化



    def choose_model_path(self):
        # 选择本地YOLO模型文件路径
        self.model_path = filedialog.askopenfilename(title="选择YOLO模型", filetypes=[("PyTorch模型", "*.pt")])
        print(f"选择的YOLO模型路径: {self.model_path}")

    def choose_save_path(self):
        # 选择保存图像的路径
        self.save_path = filedialog.askdirectory()
        print(f"选择的保存路径: {self.save_path}")

    def start_detection(self):
        if not self.model_path:
            print("请先选择YOLO模型文件")
            return

        # 加载选择的YOLO模型
        self.model = YOLO(self.model_path)

        # 打开摄像头
        cap = cv2.VideoCapture(0)

        # 设置时间和检测相关参数
        last_saved_time = time.time()

        while True:
            ret, frame = cap.read()
            if not ret:
                break

            # 人脸检测
            results = self.model(frame)
            detections = results[0].boxes  # 获取检测结果

            # 处理人脸检测框
            for detection in detections:
                x1, y1, x2, y2 = detection.xyxy[0]  # 获取检测框的坐标（左上角和右下角）
                conf = detection.conf[0]  # 获取置信度
                cls = int(detection.cls[0])  # 获取类别

                if cls == 0:  # 0 类别代表人脸
                    cv2.rectangle(frame, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2)
                    self.face_detected = True

            # 检测图像变化
            current_time = time.time()
            if self.prev_frame is not None and current_time - last_saved_time >= self.frame_rate:
                frame_diff = cv2.absdiff(frame, self.prev_frame)
                gray_diff = cv2.cvtColor(frame_diff, cv2.COLOR_BGR2GRAY)
                _, thresh = cv2.threshold(gray_diff, 50, 255, cv2.THRESH_BINARY)
                diff_area = np.sum(thresh) / 255

                # 如果变化区域足够大并且有检测到人脸
                if diff_area > 5000 and self.face_detected:
                    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
                    if self.save_path:
                        save_filename = f"{self.save_path}/face_{timestamp}.jpg"
                        cv2.imwrite(save_filename, frame)
                        print(f"图像保存为: {save_filename}")
                    last_saved_time = current_time

            # 更新当前帧
            self.prev_frame = frame.copy()

            # 显示检测结果
            #cv2.imshow("YOLO Face Detection", frame)

            # 按键检测，按'q'退出
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break

        cap.release()
        cv2.destroyAllWindows()


if __name__ == "__main__":
    root = tk.Tk()
    app = FaceDetectionApp(root)
    root.mainloop()
